package com.chengxiaoxiao.ai_demo.config;

import cn.hutool.core.collection.CollUtil;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.chengxiaoxiao.ai_demo.entity.MessageInfo;
import com.chengxiaoxiao.ai_demo.mapper.MessageInfoMapper;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import jakarta.annotation.Resource;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.List;

import static dev.langchain4j.data.message.ChatMessageDeserializer.messageFromJson;
import static dev.langchain4j.data.message.ChatMessageSerializer.messageToJson;

/**
 * 基于SQLite实现的对话历史记录存储
 *
 * @Description:
 * @Author Cheng Xiaoxiao
 * @Date: 2025/3/10 09:16
 */
@Service
public class MySqlChatMemoryStore implements ChatMemoryStore {
    @Resource
    private MessageInfoMapper messageInfoMapper;

    @Override
    public List<ChatMessage> getMessages(Object memoryId) {
        String memoryIdStr = memoryId.toString();
        List<MessageInfo> messageInfos = messageInfoMapper.selectList(new LambdaQueryWrapper<MessageInfo>().eq(MessageInfo::getConversationId, memoryIdStr));
        if (messageInfos == null) {
            return CollUtil.newArrayList();
        }

        //将消息转换为ChatMessage
        List<ChatMessage> chatMessages = new ArrayList<>();
        for (MessageInfo messageInfo : messageInfos) {
            chatMessages.add(messageFromJson(messageInfo.getMessage()));
        }
        return chatMessages;
    }

    @Override
    public void updateMessages(Object memoryId, List<ChatMessage> messages) {
        if (CollUtil.isEmpty(messages)) {
            return;
        }
        String json = messageToJson(messages.get(messages.size() - 1));
        //保存到数据库
        MessageInfo messageInfo = new MessageInfo();
        messageInfo.setConversationId(memoryId.toString());
        messageInfo.setMessage(json);
        messageInfoMapper.insert(messageInfo);
    }

    @Override
    public void deleteMessages(Object memoryId) {
        messageInfoMapper.delete(new LambdaQueryWrapper<MessageInfo>().eq(MessageInfo::getConversationId, memoryId.toString()));
    }
}
